from tipsz8 import saveLocalData, loadLocalData, clear_cache, getFileUTime
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import math
from embedding import textEmbedding
import json
data = loadLocalData("shop.json") or []

def format_entry(entry):
    fields = [
        entry.get("ID", ""),
        entry.get("商品名称", ""),
        entry.get("用途", ""),
        entry.get("游戏", ""),
        entry.get("CPU", ""),
        entry.get("cpu类型", ""),
        entry.get("gpu类型", ""),
        entry.get("keywords", ""),
        entry.get("机箱", ""),
        entry.get("主板", ""),
        entry.get("散热器", ""),
        entry.get("电源", ""),
        entry.get("显卡", ""),
        entry.get("内存", ""),
        entry.get("硬盘", "")
    ]

    return "；".join(fields)

# 5. 分批处理并构建索引
batch_size = 30
num_batches = math.ceil(len(data) / batch_size)
print(num_batches)
metadata = {}


# 获取维度并初始化 FAISS 索引
sample = textEmbedding([format_entry(data[0])])
embedding_dim = sample.shape[1]
index = faiss.IndexFlatL2(embedding_dim)
print(sample)
for i in range(num_batches):
    batch_data = data[i * batch_size:(i + 1) * batch_size]
    texts = [format_entry(entry) for entry in batch_data]
    embeddings = textEmbedding(texts)
    
    # 添加向量到 index
    index.add(embeddings)

    # 记录元数据（ID 可自行设置为 entry["ID"]）
    for j, entry in enumerate(batch_data):
        metadata[i * batch_size + j] = entry

# 6. 保存向量索引和元信息
faiss.write_index(index, 'pc_config.index')
with open("pc_config_meta.json", "w", encoding="utf-8") as f:
    json.dump(metadata, f, ensure_ascii=False, indent=2)